import math

import numpy as np
import pandas as pd

from Index.dataset import DataSet
from DataControl.getDataFromTushare.tudata import pullDataFromTushare

"""
做数据的前期准备，包括returns，direction，lag_1, lag_2, lag_1_bin, lag_2_bin
"""
class DNN_Dataset(object):
    def __init__(self, ts_code, start_date, end_date):
        self.ts_code = ts_code
        self.data = DataSet(ts_code).data
        self.lags = 2
       # print(self.data.shape)
       # print(self.data.columns)
        self.data['returns'] = np.zeros(self.data.shape[0])
        pre = self.data.index[0]
        for i in self.data.index[1:]:
            self.data.loc[i, 'returns'] = math.log(self.data.loc[i, 'close'] / self.data.loc[pre, 'close'])
            pre = i
        self.data['direction'] = np.sign(self.data['returns']).astype(int)

        self.cols = self.create_lags()
        self.data.dropna(inplace = True)

        self.cols_bin = self.create_bins()

    # lag_1是前一天的returns， lag_2是前天的returns
    def create_lags(self):
        cols = []
        for lag in range(1, self.lags + 1):
            col = 'lag_{}'.format(lag)
            self.data[col] = self.data['returns'].shift(lag)
            cols.append(col)
        return cols

    # bin是lag的二值化，lag小于0为-1， lag大于0为1
    def create_bins(self, bins = [0]):
        cols_bin = []
        for col in self.cols:
            col_bin = col + '_bin'
            self.data[col_bin] = np.digitize(self.data[col], bins=bins)
            cols_bin.append(col_bin)
        return cols_bin
if __name__ == '__main__':
    data = DNN_Dataset('000001.SZ')

